# -*- coding: utf-8 -*- 
# Author: Jacky
# Creation Date: 2021/3/24

import numpy as np
import torch
from collections import Counter
from torch import nn
from torch.utils.data import DataLoader, Dataset
from data_set import convert
'''当我们得到一个数据集时，Dataset类可以帮我们提取我们需要的数据，我们用子类继承Dataset类，我们先给每个数据一个编号（idx），在后面的神经网络中，初始化Dataset子类实例后，就可以通过这个编号去实例对象中读取相应的数据，会自动调用__getitem__方法，同时子类对象也会获取相应真实的Label（人为去复写即可）'''

# 用于加载电影评分数据集
class MovieDataset(Dataset):
    """An abstract Dataset class wrapped around Pytorch Dataset class.
        这里假设，输入的数据集为sample 的矩阵形式
    """

    def __init__(self, dataset, idxs): # 输入的数据集和数据集中选择的索引
        self.dataset = dataset
        self.idxs = [int(i) for i in idxs] # 将输入的索引列表转换为整数类型，并保存在self.idxs属性中
        # self.user_movie = convert(self.dataset[self.idxs], max(self.dataset[:, 1]))

    def __len__(self): # 获取数据集的长度
        return len(self.idxs)

    def __getitem__(self, item): # 用于获取数据集中特定索引位置的数据
        ratings = self.dataset[self.idxs[item], 1:-3] # 获取从第一列到倒数第四个列之间，在idxs索引行之间的数据，即根据user_id排序之后user、movie_id、rating
        return torch.tensor(ratings) # 将获取的数据转换为PyTorch的Tensor类型


# 用于在每个客户端上训练本地模型。在初始化方法中，将输入的数据集转换为用户-电影评分矩阵，并创建数据加载器
class LocalUpdate(object):
    def __init__(self, args, dataset, idxs): # idxs决定选择的batch,dataset为sample但只选择1-3列数据
        self.args = args
        # 将输入的限制索引的dataset转换为user_movie
        self.dataset = convert(dataset[idxs], int(max(dataset[:, 1]))) # 此处dataset即为data_set = np.array(sample)，第二列（索引为1）的最大值即为最大电影数，self.dataset为更新后的二维数组
        self.idxs = np.arange(0, len(self.dataset))
        # 创建一个 DataLoader 对象，用于加载 MovieDataset 类型的数据集。MovieDataset 是一个自定义的 PyTorch 数据集类，用于处理影片数据。这个 DataLoader 将数据集按照指定的 batch_size 大小和 shuffle 参数进行分批和混洗。
        # shuffle=True 表示在每个 epoch 开始之前对数据进行混洗（随机重排）。这意味着在每个训练周期开始时，DataLoader 会重新随机排列数据集中的样本，以防止模型在训练过程中对数据顺序的依赖性。
        self.trainloader = DataLoader(MovieDataset(self.dataset, self.idxs),
                                      batch_size=self.args.local_bs, shuffle=True) # args.local_bs即为local batch size，默认为50

        # self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
        self.device = 'cuda' if args.gpu else 'cpu'
        # Default criterion set to MSE Loss function
        # 创建了一个均方误差损失函数对象，并将其移动到指定的设备上，以便后续在该设备上计算损失
        self.criterion = nn.MSELoss().to(self.device)

    # 训练本地模型并返回模型参数和训练损失，global_round表示当前回合数，对应FPCC 75行，唯一用法为model对应AFPC中的global_model，global_mode由SAE获得，因为是SAE(神经网络)，所以输出值和输入值ratings本身比较
    def update_weights(self, model, client_idx, global_round): # model传入参数为全局模型
        """
        训练本地模型，得到模型参数和训练loss
        :param model:
        :param client_idx: 客户0~9
        :param global_round: 全局回合数
        :return: model.state_dict() 模型参数
        :return: sum(epoch_loss) / len(epoch_loss) 本地训练损失
        """
        # Set mode to train model
        # 指定为train模式
        model.train()
        epoch_loss = []

        # 根据参数设置选择优化器（SGD 或 Adam）
        if self.args.optimizer == 'sgd':
            # .parameters()是 PyTorch 中 nn.Module 类的一个方法，用于获取模型中的所有参数（可学习的权重和偏置）。这个方法返回一个生成器，它会逐层遍历模型，并将每一层的参数返回。
            # lr：学习率，momentum：动量
            optimizer = torch.optim.SGD(model.parameters(), lr=self.args.lr,
                                        momentum=0.5)
        elif self.args.optimizer == 'adam':
            # weight_decay：权重衰减
            optimizer = torch.optim.Adam(model.parameters(), lr=self.args.lr,
                                         weight_decay=1e-4)
        # 本地训练 训练回合数设置为local_ep
        for iter in range(self.args.local_epoch): # 对于每个本地训练回合
            batch_loss = []
            # 在一个循环中，用于遍历 self.trainloader 中的每个批次数据。在每个迭代中，enumerate 函数会返回一个迭代器，其中包含批次的索引 batch_idx 和对应的数据 ratings。
            # enumerate() 函数用于将一个可遍历的数据对象(如列表、元组或字符串)组合为一个索引序列，同时列出数据和数据下标，一般用在 for 循环当中。
            for batch_idx, ratings in enumerate(self.trainloader):
                ratings = ratings.to(self.device) # 将数据 ratings 移动到指定的设备上，通常是 GPU 或 CPU

                model.zero_grad() # 将模型参数的梯度归零。
                outputs = model(ratings) # 将数据传递给模型并获取输出。
                loss = self.criterion(outputs, ratings) # 计算输出与真实标签之间的损失，使用均方误差损失函数criterion
                loss.backward() # 根据损失计算模型参数的梯度。
                optimizer.step() # 使用优化器更新模型参数。
                # 控制是否输出训练信息，self.args.verbose 是一个布尔值，用于指示是否启用详细输出。batch_idx % 5 == 0 表示每训练 5 个批次输出一次信息，以减少输出的频率。
                if self.args.verbose and (batch_idx % 5 == 0):
                    print(
                        '| Global Round : {} | Client : {} | Local Epoch : {} | [{}/{} ({:.0f}%)]\tLoss: {:.6f}'.format(
                            global_round, client_idx, iter + 1, batch_idx * len(ratings),
                            len(self.trainloader.dataset), 100. * batch_idx / len(self.trainloader), loss.item()))
                # 将每个批次的损失值 loss.item() 添加到列表 batch_loss 中
                batch_loss.append(loss.item())
                # print(model.state_dict()['linear.weight'].numpy())

            epoch_loss.append(sum(batch_loss) / len(batch_loss))
        # 返回第一个值是当前模型的状态字典（权重张量和偏置张量字典），即模型的所有参数和缓冲区的字典。第二个值是整个本地训练过程中损失值的平均值，即每个训练周期的平均损失
        return model.state_dict(), sum(epoch_loss) / len(epoch_loss)


# 用于在每个客户端上训练本地模型，与LocalUpdate类相似，但具有额外的first_update_weights方法，用于第一次更新权重
class DQN_LocalUpdate(object):
    def __init__(self, args, dataset, idxs, local_epoch): # dataset也是sample，相比于LocalUpdate唯一多了一个local_epoch
        self.args = args
        self.local_epoch = local_epoch
        # 将输入的dataset转换为user_movie
        self.dataset = convert(dataset[idxs], int(max(dataset[:, 1])))
        # A = (self.distribution[:,0]).tolist()
        # B = (self.distribution[:,1]).tolist()
        # C = [A, B]
        self.idxs = np.arange(0, len(self.dataset))
        self.trainloader = DataLoader(MovieDataset(self.dataset, self.idxs),
                                      batch_size=self.args.local_bs, shuffle=True)
        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
        # self.device = 'cuda' if args.gpu else 'cpu'
        # Default criterion set to MSE Loss function
        self.criterion = nn.MSELoss().to(self.device)

    # 与update_weights结构上一致，对应AFPC 72行
    def first_update_weights(self, model, client_idx, global_round): #model为全局模型
        """
        训练本地模型，得到模型参数和训练loss
        :param model:
        :param client_idx: 客户0~9
        :param global_round: 全局回合数
        :return: model.state_dict() 模型参数
        :return: sum(epoch_loss) / len(epoch_loss) 本地训练损失
        """
        # Set mode to train model
        model.train()
        epoch_loss = []

        # Set optimizer for the local updates
        if self.args.optimizer == 'sgd':
            optimizer = torch.optim.SGD(model.parameters(), lr=self.args.lr,
                                        momentum=0.5)
        elif self.args.optimizer == 'adam':
            optimizer = torch.optim.Adam(model.parameters(), lr=self.args.lr,
                                         weight_decay=1e-4)
        # 本地训练 训练回合数设置为local_ep
        for iter in range(self.local_epoch):
            batch_loss = []
            for batch_idx, ratings in enumerate(self.trainloader):
                ratings = ratings.to(self.device)

                model.zero_grad()
                outputs = model(ratings)
                loss = self.criterion(outputs, ratings)
                loss.backward()
                optimizer.step()

                if self.args.verbose and (batch_idx % 5 == 0):
                    print(
                        '| Global Round : {} | Client : {} | Local Epoch : {} | [{}/{} ({:.0f}%)]\tLoss: {:.6f}'.format(
                            global_round, client_idx, iter + 1, batch_idx * len(ratings),
                            len(self.trainloader.dataset), 100. * batch_idx / len(self.trainloader), loss.item()))
                batch_loss.append(loss.item())
                # print(model.state_dict()['linear.weight'].numpy())

            epoch_loss.append(sum(batch_loss) / len(batch_loss))
        return model.state_dict(), sum(epoch_loss) / len(epoch_loss)

    # 与update_weights结构上一致，用于第一次更新权重，对应AFPC 118行
    def update_weights(self, model, client_idx, global_round):
        """
        训练本地模型，得到模型参数和训练loss
        :param model:
        :param client_idx: 客户0~9
        :param global_round: 全局回合数
        :return: model.state_dict() 模型参数
        :return: sum(epoch_loss) / len(epoch_loss) 本地训练损失
        """
        # Set mode to train model
        model.train()
        epoch_loss = []

        # Set optimizer for the local updates
        if self.args.optimizer == 'sgd':
            optimizer = torch.optim.SGD(model.parameters(), lr=self.args.lr,
                                        momentum=0.5)
        elif self.args.optimizer == 'adam':
            optimizer = torch.optim.Adam(model.parameters(), lr=self.args.lr,
                                         weight_decay=1e-4)
        # 本地训练 训练回合数设置为local_ep
        for iter in range(self.local_epoch):
            batch_loss = []
            for batch_idx, ratings in enumerate(self.trainloader):
                ratings = ratings.to(self.device)

                model.zero_grad()
                outputs = model(ratings)
                loss = self.criterion(outputs, ratings)
                loss.backward()
                optimizer.step()

                if self.args.verbose and (batch_idx % 5 == 0):
                    print(
                        '| Global Round : {} | Client : {} | Local Epoch : {} | [{}/{} ({:.0f}%)]\tLoss: {:.6f}'.format(
                            global_round, client_idx, iter + 1, batch_idx * len(ratings),
                            len(self.trainloader.dataset), 100. * batch_idx / len(self.trainloader), loss.item()))
                batch_loss.append(loss.item())
                # print(model.state_dict()['linear.weight'].numpy())

            epoch_loss.append(sum(batch_loss) / len(batch_loss))
        return model.state_dict(), sum(epoch_loss) / len(epoch_loss)


# 用于计算缓存命中率
def cache_hit_ratio(test_dataset, cache_items):
    """
    计算缓存命中率
    :param test_dataset: user_group_test[0-9] dataset
    :param cache_items: 缓存内容
    :return: 缓存命中率
    """
    requset_items = test_dataset[:, 1] # 所有行第二列数据，返回一维数组
    count = Counter(requset_items) # 统计 requset_items 中每个元素出现的次数
    CACHE_HIT_NUM = 0
    for item in cache_items:
        CACHE_HIT_NUM += count[item]
    CACHE_HIT_RATIO = CACHE_HIT_NUM / len(requset_items) * 100
    return CACHE_HIT_RATIO

# 将数据列表写入到指定路径的 CSV 文件中
def text_save(filename, data):#filename为写入CSV文件的路径，data为要写入数据列表.
    file = open(filename,'a')
    for i in range(len(data)):
        s = str(data[i]).replace('[','').replace(']','')#去除[],这两行按数据不同，可以选择
        s = s.replace("'",'').replace(',','') +'\n'   #去除单引号，逗号，每行末尾追加换行符
        file.write(s)
    file.close()
    print("保存成功")





